import datetime
import numpy as np
from scipy.interpolate import griddata
from tqdm import tqdm  # 添加在文件开头

import dask.array as da  # 用于分片并行计算
from dask import delayed  # 用于延迟执行插值函数
import time

def readloc(filename, time_col=0, lat_col=1, lon_col=2, depth_col=3):
    """
    读取导航数据的时间戳及对应空间坐标（lon,lat,depth）
    :param filename: 文件名
    :return: 时间戳列表和对应的空间坐标列表
    """
    with open(filename, 'r') as file:
        lines = file.readlines()

    timestamps = []
    coordinates = []

    for line in lines:
        parts = line.strip().split(',')
        if len(parts) >= 3:
            timestamp_str = parts[time_col].strip()
            timestamp = datetime.datetime.strptime(timestamp_str, '%Y-%m-%dT%H:%M:%S.%f')
            lon = float(parts[lon_col].strip())
            lat = float(parts[lat_col].strip())
            depth = float(parts[depth_col].strip())
            
            timestamps.append(timestamp)
            coordinates.append((lon, lat, depth))
    return timestamps, coordinates

def match_coordinates(timestamps, coordinates, data_times):
    """
    将时间戳与坐标数据匹配
    :param timestamps: 时间戳列表
    :param coordinates: 坐标列表
    :param data_times: 数据时间列表
    :return: 匹配后的坐标列表
    """
    matched_coordinates = []
    start_i = 0
    for data_time in data_times:
        #根据相邻时间的坐标内插
        
        for i in range(start_i,len(timestamps) - 1):
            if timestamps[i] <= data_time <= timestamps[i + 1]:
                
                if timestamps[i] == data_time:
                    #如果时间戳完全匹配，直接使用对应坐标
                    matched_coordinates.append(coordinates[i])
                    start_i = i
                    break
                elif timestamps[i + 1] == data_time:
                    #如果时间戳完全匹配，直接使用对应坐标
                    matched_coordinates.append(coordinates[i + 1])
                    start_i = i
                    break
                else:
                    # 线性插值
                    t1 = timestamps[i]
                    t2 = timestamps[i + 1]
                    ratio = (data_time - t1) / (t2 - t1)
                    
                    lon = coordinates[i][0] + ratio * (coordinates[i + 1][0] - coordinates[i][0])
                    lat = coordinates[i][1] + ratio * (coordinates[i + 1][1] - coordinates[i][1])
                    depth = coordinates[i][2] + ratio * (coordinates[i + 1][2] - coordinates[i][2])
                    
                    matched_coordinates.append((lon, lat, depth))
                    start_i = i
                    break
        # print(data_time, "matched with coordinates:", matched_coordinates[-1])
    return matched_coordinates

def SEC2(data, sample_rate=0.4,vel=0.0333,attenuation=5,start_gain=1.0,maximum_gain=500.0):
    """
    对雷达数据进行增益处理
    :param data: 雷达数据 (sample_num, trace_num)
    :param sample_rate: 采样间隔（纳秒）
    :param vel: 电磁波速度（米/纳秒）
    :param attenuation: 衰减系数（dB/m）
    :param start_gain: 初始增益
    :param maximum_gain: 最大增益
    :return: 增益后的数据
    """
    sample_num, trace_num = data.shape
    depths = (np.arange(sample_num)+1) * sample_rate * vel / 2
    gains = []
    for depth in depths:
        gain = start_gain * np.exp(attenuation * depth / 20)  # 计算增益
        gain = min(gain, maximum_gain)  # 限制最大增益
        gains.append(gain)
    gains = np.array(gains).reshape(-1, 1)  # 转换为列向量(53,1)
    # 使用np.multiply确保正确广播
    agc_data = np.multiply(data, gains)  # 将(53,1)广播到(53,14221)
    return agc_data

def AGC(data, window_size):
    """
    对雷达数据应用AGC增益
    :param data: 雷达数据
    :param window_size: 增益统计时窗
    :return: 增益后的数据
    """
    sample_num, trace_num = data.shape
    agc_data = np.zeros_like(data, dtype=np.float64)  # 使用 float64 避免溢出

    for j in range(trace_num):
        # print("processing trace num: %d"%(j+1))
        for i in range(sample_num):
            start = max(0, i - window_size // 2)
            end = min(sample_num, i + window_size // 2 + 1)
            
            # 计算局部均方根
            window = data[start:end, j]
            win2 = []
            for win in window:
                win2.append(win*win)
            # print(win2)
            rms = np.sqrt(np.mean(win2))
            
            # 避免除以零
            if rms > 0:
                gain = 1 / rms
            else:
                gain = 1.0  # 如果 rms 为零，则增益为 1

            agc_data[i, j] = data[i, j] * gain # 应用增益
            
    return agc_data


def data_interpolate(data, matched_coordinates, d_horizontal, dz=0.001, vel=0.0333, sample_rate=0.4):
    """
    对雷达数据进行插值
    :param data: 雷达数据 (sample_num, trace_num)
    :param matched_coordinates: 坐标数据 (trace_num, 3) [lon, lat, depth]
    :param d_horizontal: 水平插值间隔（度）
    :param dz: 深度插值间隔（米）
    :param vel: 电磁波速度（米/纳秒）
    :param sample_rate: 采样间隔（纳秒）
    :return: 插值后的lon, lat, depth网格，以及插值后的数据 (grid_lon, grid_lat, grid_depth, grid_data)
    """

    # 检查输入数据的维度
    # 构造所有点的三维坐标和数值
    sample_num, trace_num = data.shape
    lon = np.array([coord[0] for coord in matched_coordinates])
    lat = np.array([coord[1] for coord in matched_coordinates])
    depth = np.array([coord[2] for coord in matched_coordinates])

    # 每个采样点的实际深度
    depth_axis = depth.mean() + np.arange(sample_num) * sample_rate * vel / 2
    # depth_axis 保留5位小数
    depth_axis = np.round(depth_axis, 5)
    # 构造三维点和数值
    points = []
    values = []
    for i in range(sample_num):
        for j in range(trace_num):
            points.append([lon[j], lat[j], depth_axis[i]])
            values.append(data[i, j])
    points = np.array(points)
    values = np.array(values)

    # 构造三维网格
    grid_lon = np.arange(lon.min(), lon.max(), d_horizontal)
    grid_lat = np.arange(lat.min(), lat.max(), d_horizontal)
    grid_depth = np.arange(depth_axis.min(), depth_axis.max(), dz)
    grid_lon, grid_lat, grid_depth = np.meshgrid(grid_lon, grid_lat, grid_depth, indexing='ij')

    # 三维插值
    grid_data = griddata(points, values, (grid_lon, grid_lat, grid_depth), method='linear')

    return grid_lon, grid_lat, grid_depth, grid_data

def data_interpolate(data, matched_coordinates, d_horizontal, dz=0.001, vel=0.0333, sample_rate=0.4):
    """
    对雷达数据进行插值（多线程版，基于Dask分片并行，不修改原有插值逻辑）
    :param data: 雷达数据 (sample_num, trace_num)
    :param matched_coordinates: 坐标数据 (trace_num, 3) [lon, lat, depth]
    :param d_horizontal: 水平插值间隔（度）
    :param dz: 深度插值间隔（米）
    :param vel: 电磁波速度（米/纳秒）
    :param sample_rate: 采样间隔（纳秒）
    :return: 插值后的lon, lat, depth网格，以及插值后的数据 (grid_lon, grid_lat, grid_depth, grid_data)
    """

    # --------------------------
    # 1. 原有逻辑：构造三维离散点和数值（完全不变）
    # --------------------------
    sample_num, trace_num = data.shape
    lon = np.array([coord[0] for coord in matched_coordinates])
    lat = np.array([coord[1] for coord in matched_coordinates])
    depth = np.array([coord[2] for coord in matched_coordinates])

    # 每个采样点的实际深度
    depth_axis = depth.mean() + np.arange(sample_num) * sample_rate * vel / 2
    depth_axis = np.round(depth_axis, 5)  # 保留5位小数

    # 构造三维点和数值（原有双层循环逻辑不变）
    points = []
    values = []
    for i in range(sample_num):
        for j in range(trace_num):
            points.append([lon[j], lat[j], depth_axis[i]])
            values.append(data[i, j])
    points = np.array(points)
    values = np.array(values)

    # --------------------------
    # 2. 原有逻辑：生成规则三维网格（完全不变）
    # --------------------------
    grid_lon = np.arange(lon.min(), lon.max(), d_horizontal)
    grid_lat = np.arange(lat.min(), lat.max(), d_horizontal)
    grid_depth = np.arange(depth_axis.min(), depth_axis.max(), dz)
    # 生成网格（indexing='ij'保持原有格式）
    grid_lon, grid_lat, grid_depth = np.meshgrid(
        grid_lon, grid_lat, grid_depth, indexing='ij'
    )

    # --------------------------
    # 3. 新增：多线程改造（分片并行插值）
    # --------------------------
    def _interp_chunk(chunk_lon, chunk_lat, chunk_depth):
        """
        单个分片的插值函数（内部完全复用原有griddata逻辑）
        注意：此函数会被每个分片调用，无需修改
        """
        return griddata(
            points=points,
            values=values,
            xi=(chunk_lon, chunk_lat, chunk_depth),
            method='linear'  # 保留原有线性插值
        )

    # 3.1 将NumPy网格数组转换为Dask数组，并分片（核心：多线程的基础）
    # 分片大小：根据CPU核心数调整（示例：lon方向分4片，lat方向分4片，depth方向不分片）
    # 原则：分片总数 ≈ CPU核心数（避免分片过多/过少）
    chunk_size = (
        max(1, grid_lon.shape[0] // 4),  # lon方向分片数（4片，可根据CPU核心数调整）
        max(1, grid_lat.shape[1] // 4),   # lat方向分片数（4片）
        grid_depth.shape[2]               # depth方向不分片（保持连续性）
    )

    # 转换为Dask数组（自动支持多线程）
    dask_lon = da.from_array(grid_lon, chunks=chunk_size)
    dask_lat = da.from_array(grid_lat, chunks=chunk_size)
    dask_depth = da.from_array(grid_depth, chunks=chunk_size)

    # 3.2 并行执行插值（对每个分片调用_interp_chunk，自动多线程）
    # map_blocks：将函数应用到每个Dask分片
    dask_grid_data = da.map_blocks(
        _interp_chunk,  # 单个分片的插值函数（复用原有逻辑）
        dask_lon, dask_lat, dask_depth,
        meta=np.array([], dtype=np.float64)  # 声明输出数据类型（匹配原有griddata结果）
    )

    # 3.3 触发计算并转换回NumPy数组（保持原有返回格式）
    # compute()：执行并行计算，默认用多线程（可通过scheduler参数指定线程数）
    grid_data = dask_grid_data.compute(scheduler='threads', num_workers=8)  # num_workers=CPU核心数

    # --------------------------
    # 4. 原有逻辑：返回结果（完全不变）
    # --------------------------
    return grid_lon, grid_lat, grid_depth, grid_data

def generate_slices(data, matched_coordinates, d_horizontal_slice=0.0000001, vel=0.0333, sample_rate=0.4,
                    slice_interval=10, slice_thickness=7, neighborhood_radius=0.5):
    """
    生成切片数据
    :param data: 雷达数据 (sample_num, trace_num)
    :param matched_coordinates: 坐标数据 (trace_num, 3) [lon, lat, depth]
    :param d_horizontal_slice: 切片插值水平间隔（度）
    :param slice_interval: 切片间隔（个采样点）
    :param slice_thickness: 切片厚度（个采样点，应为奇数）
    :return: 切片的经度、纬度、深度和数值
    """
    sample_num, trace_num = data.shape
    # 构造切片的经度和纬度
    lon = np.array([coord[0] for coord in matched_coordinates])
    lat = np.array([coord[1] for coord in matched_coordinates])
    depth = np.array([coord[2] for coord in matched_coordinates])
    points = np.column_stack((lon, lat))  # shape: (trace_num, 2)
    # print(points)
    # 每个采样点的实际深度
    depth_axis = depth.mean() + np.arange(sample_num) * sample_rate * vel / 2
    depth_axis = np.round(depth_axis, 5)

    # 构造二维网格点
    grid_lon = np.arange(lon.min(), lon.max(), d_horizontal_slice)
    grid_lat = np.arange(lat.min(), lat.max(), d_horizontal_slice)
    grid_lon, grid_lat = np.meshgrid(grid_lon, grid_lat, indexing='ij')

    depths = []
    weighted_average_grid_datas = []  # 存储每个切片的加权平均值
    # 如果 slice_thickness 不是奇数，则减1
    if slice_thickness % 2 == 0:
        slice_thickness -= 1
    half_thickness = slice_thickness // 2

    # 计算进度条总数
    total_steps = (sample_num - 2 * half_thickness) // slice_interval

    for idx, i in enumerate(tqdm(range(half_thickness, sample_num - half_thickness, slice_interval), total=total_steps,
                                 desc="切片处理中")):
        depth = depth_axis[i]
        # print(f"切片处理中 {idx+1}/{total_steps}")  # 可去掉原print

        # 取切片数据并计算平均值
        slice_data = data[i - half_thickness:i + half_thickness + 1, :]  # shape: (thickness, trace_num)
        slice_data_average = np.mean(slice_data, axis=0)  # shape: (trace_num,)
        slice_data_average = slice_data_average.reshape(1, -1)  # 转换为二维数组以便后续处理
        slice_data_average_flat = []  # 用于存储展平后的数据

        # 检查数据维度
        for j in range(slice_data_average.shape[1]):
            slice_data_average_flat.append(slice_data_average[0, j])
        slice_data_average_flat = np.array(slice_data_average_flat)
        # 使用半径为Neighborhood_Radius米内的数据，并按距离加权计算格点值
        weighted_average_grid_data = weighted_average_grid(points, slice_data_average_flat, grid_lon, grid_lat,
                                                           radius=neighborhood_radius)
        # 插值数据
        # grid_data = griddata(points, slice_data_average_flat, (grid_lon, grid_lat), method='linear')
        depths.append(depth)
        weighted_average_grid_datas.append(weighted_average_grid_data)
    return lon.min(), lon.max(), lat.min(), lat.max(), weighted_average_grid_datas, depths

def weighted_average_grid(points, values, grid_lon, grid_lat, radius=0.0001, eps=1e-12):
    """
    对每个格点，查找半径内的点，按距离加权平均
    points: (N,2) 采样点经纬度
    values: (N,) 采样点数值
    grid_lon, grid_lat: meshgrid后的格点
    radius: 半径（m）
    eps: 防止除零
    返回：与grid_lon同shape的加权平均值
    """
    # 将radius的m转换为度
    radius = radius / 111320  # 1度约等于111320米
    grid_shape = grid_lon.shape
    grid_value = np.full(grid_shape, np.nan)
    points = np.asarray(points)
    values = np.asarray(values)
    for i in range(grid_shape[0]):
        for j in range(grid_shape[1]):
            gx, gy = grid_lon[i, j], grid_lat[i, j]
            dists = np.sqrt((points[:, 0] - gx) ** 2 + (points[:, 1] - gy) ** 2)
            mask = dists <= radius
            if np.any(mask):
                w = 1 / (dists[mask] + eps)
                grid_value[i, j] = np.sum(w * values[mask]) / np.sum(w)
    return grid_value

from scipy.signal import butter, filtfilt

def bandpass_filter(data, lowcut, highcut, fs, order=5):
    nyquist = 0.5 * fs
    low = lowcut / nyquist
    high = highcut / nyquist
    b, a = butter(order, [low, high], btype='band')
    filtered_data = filtfilt(b, a, data, axis=0)
    return filtered_data